from typing import List, Dict, Any
from rainforeLearn.gomoku.v2.train.constants.train_constants import TrainConstants


class BatchTracker:
    """批次统计跟踪器 - 负责批次相关的统计数据"""

    def __init__(self):
        self.win_rates: List[float] = []
        self.batch_game_results: List[int] = []

    def add_batch_game_result(self, winner: int) -> None:
        """添加游戏结果"""
        self.batch_game_results.append(winner)

    def calculate_current_win_rate(self) -> float:
        """计算当前批次胜率"""
        if not self.batch_game_results:
            return 0.0

        ai_wins = sum(1 for result in self.batch_game_results if result == TrainConstants.AI_PLAYER)
        return ai_wins / len(self.batch_game_results)

    def complete_batch(self) -> float:
        """完成当前批次，返回胜率"""
        current_win_rate = self.calculate_current_win_rate()
        self.win_rates.append(current_win_rate)

        # 重置当前批次
        self.batch_game_results = []
        return current_win_rate
